package comlib

import chisel3._
import chisel3.util._
import comlib.util.{muxO}

class Crossbar[D<:Data](data:D, inNum:Int, outNum:Int, inFF:Boolean=false, selFF:Boolean=false, nSel:Int=1, outFF:Boolean=false) extends MultiIOModule{
  
  require(nSel>0, f"分段必须为自然数")
  
  val in = IO(Input(Vec(inNum, chiselTypeOf(data))))
  val ilock = IO(Input(Bool()))
  val olock = IO(Input(Bool()))
  val sel = IO(Input(Vec(nSel, Vec(outNum, Vec(inNum, Bool())))))
  val out = IO(Output(Vec(outNum, chiselTypeOf(data))))
  val data_in_ff = IO(Output(Vec(inNum, chiselTypeOf(data))))
  val sel_ff = IO(Output(Vec(outNum, Vec(inNum, Bool()))))
  val out_pre_ff = IO(Output(Vec(outNum, chiselTypeOf(data))))
  val out_valid = IO(Output(Vec(outNum, Bool())))

  val in_ff1 = if(inFF) RegEnable(in, !ilock) else in
  val in_inner = in_ff1.map(_.asUInt)
  val sel_inner = if( selFF ) RegEnable(sel, !ilock) else sel

  val res_inner = if( nSel > 1 ){
    val width = in_inner(0).getWidth
    val fullDataSize = scala.math.ceil(width / nSel).toInt
    val rmnDataSize = width - fullDataSize * (nSel-1)

    val res = (0 until outNum)
      .map { case o => 
        val fullData = (nSel-2 to 0 by -1).map { case s =>
          val mux_lsb = fullDataSize * s + rmnDataSize
          val mux_msb = mux_lsb + fullDataSize - 1
          val mux_in = in_inner.map { case d => d(mux_msb, mux_lsb) }
          muxO( VecInit(mux_in), sel_inner(s)(o) )
        }
        val rmnData = {
          val mux_in = in_inner.map { case d => d(rmnDataSize-1, 0) }
          muxO( VecInit(mux_in), sel_inner(nSel-1)(o) )
        }
        Cat( fullData :+ rmnData )
      }

    VecInit(res)
  }else{
    VecInit( sel_inner(0).map { case s => muxO( VecInit(in_inner), s ) } )
  }

  val out_inner = VecInit(res_inner.map(_.asTypeOf(data)))

  if( outFF ) out := RegEnable(out_inner, !olock)
  else out := out_inner

  val sel_out_valid = VecInit(sel_inner.map{case s => s.reduce(_|_)})}
  out_valid := Mux( outFF.B, RegNext( sel_out_valid ), sel_out_valid )

  data_in_ff := in_ff1
  sel_ff := sel_inner(0)
  out_pre_ff := out_inner
}
case class CrossbarCfg(inFF:Boolean=false, selFF:Boolean=false, nSel:Int=1, outFF:Boolean=false)
object Crossbar{

  def apply[D<:Data](in:Vec[D], sel: => Vec[Vec[Vec[Bool]]], cfg:CrossbarCfg):Crossbar[D] = {
    val xb = Module( new Crossbar(in(0), in.size, sel.size, cfg.inFF, cfg.selFF, sel.size, cfg.outFF) )
    xb.in := in 
    xb.sel := sel 
    xb.ilock := false.B
    xb.olock := false.B
    xb
  }

  def apply[D<:Data](in:Vec[D], sel:Vec[Vec[Bool]], inFF:Boolean=false, selFF:Int=0, outFF:Boolean=false):Crossbar[D] = {
    val selFF_i = selFF > 0
    val nSel_i = if( selFF>0 ) selFF else 1
    val xb = Module( new Crossbar(in(0), in.size, sel.size, inFF, selFF_i, nSel_i, outFF) )
    xb.in := in 
    xb.sel := VecInit(Seq.fill(nSel_i)(sel)) 
    xb.ilock := false.B
    xb.olock := false.B
    xb
  }

  def xbO[D<:Data](in:Vec[D], sel:Vec[UInt], inFF:Boolean=false, selFF:Int=0, outFF:Boolean=false):Crossbar[D] = {
    val selFF_i = selFF > 0
    val nSel_i = if( selFF>0 ) selFF else 1
    val xb = Module( new Crossbar(in(0), in.size, sel.size, inFF, selFF_i, nSel_i, outFF) )
    xb.in := in 
    val selBools = VecInit(sel.map{case s => VecInit(s.asBools)})
    xb.sel := VecInit(Seq.fill(nSel_i)(selBools)) 
    xb.ilock := false.B
    xb.olock := false.B
    xb
  }

  def xbB[D<:Data](in:Vec[D], sel:Vec[UInt], inFF:Boolean=false, selFF:Int=0, outFF:Boolean=false):Crossbar[D] = {
    val selFF_i = selFF > 0
    val nSel_i = if( selFF>0 ) selFF else 1
    val xb = Module( new Crossbar(in(0), in.size, sel.size, inFF, selFF_i, nSel_i, outFF) )
    xb.in := in 
    val selBools = VecInit(sel.map{case s => VecInit(UIntToOH(s).asBools)})
    xb.sel := VecInit(Seq.fill(nSel_i)(selBools)) 
    xb.ilock := false.B
    xb.olock := false.B
    xb
  }

  def select[D<:Data](sel:Vec[Bool], in:Vec[D], outFF:Boolean=false, selFF:Int=0):Crossbar[D] = {
    val selFF_i = selFF > 0
    val nSel_i = if( selFF>0 ) selFF else 1
    val xb = Module( new Crossbar(in(0), in.size, 1, false, selFF_i, nSel_i, outFF) )
    xb.in := in 
    val selBools = VecInit(Seq.fill(1)(sel))
    xb.sel := VecInit(Seq.fill(nSel_i)(selBools)) 
    xb.ilock := false.B
    xb.olock := false.B
    xb
  }

  def selectO[D<:Data](sel:UInt, in:Vec[D], outFF:Boolean=false, selFF:Int=0):Crossbar[D] = {
    select( VecInit(sel.asBools), in, outFF, selFF )
  }

  def selectB[D<:Data](sel:UInt, in:Vec[D], outFF:Boolean=false, selFF:Int=0):Crossbar[D] = {
    selectO( UIntToOH(sel), in, outFF, selFF )
  }

  def selectB[D<:Data](sel:Seq[UInt], in:Vec[D], outFF:Boolean):Crossbar[D] = {
    val xb = Module( new Crossbar(in(0), in.size, 1, false, false, sel.size, outFF) )
    xb.in := in 
    val selBools = VecInit(sel.map{case seg => VecInit(seg.asBools)})
    xb.sel := selBools 
    xb.ilock := false.B
    xb.olock := false.B
    xb
  }

}
